Python: Conditional Average Treatment Effects (CATEs) for IRM models#

In this simple example, we illustrate how the DoubleML package can be used to estimate conditional average treatment effects with B-splines for one or two-dimensional effects in the DoubleMLIRM model.

[1]:
import numpy as np
import pandas as pd
import doubleml as dml

from doubleml.datasets import make_heterogeneous_data

Data#

We define a data generating process to create synthetic data to compare the estimates to the true effect. The data generating process is based on the Monte Carlo simulation from Oprescu et al. (2019).

The documentation of the data generating process can be found here.

One-dimensional Example#

We start with an one-dimensional effect and create our training data. In this example the true effect depends only the first covariate \(X_0\) and takes the following form

\[\theta_0(X) = \exp(2X_0) + 3\sin(4X_0).\]

The generated dictionary also contains a callable with key treatment_effect to calculate the true treatment effect for new observations.

[2]:
np.random.seed(42)
data_dict = make_heterogeneous_data(
    n_obs=2000,
    p=10,
    support_size=5,
    n_x=1,
    binary_treatment=True,
)
treatment_effect = data_dict['treatment_effect']
data = data_dict['data']
print(data.head())
          y    d       X_0       X_1       X_2       X_3       X_4       X_5  \
0  4.803300  1.0  0.259828  0.886086  0.895690  0.297287  0.229994  0.411304
1  5.655547  1.0  0.824350  0.396992  0.156317  0.737951  0.360475  0.671271
2  1.878402  0.0  0.988421  0.977280  0.793818  0.659423  0.577807  0.866102
3  6.941440  1.0  0.427486  0.330285  0.564232  0.850575  0.201528  0.934433
4  1.703049  1.0  0.016200  0.818380  0.040139  0.889913  0.991963  0.294067

        X_6       X_7       X_8       X_9
0  0.240532  0.672384  0.826065  0.673092
1  0.270644  0.081230  0.992582  0.156202
2  0.289440  0.467681  0.619390  0.411190
3  0.689088  0.823273  0.556191  0.779517
4  0.210319  0.765363  0.253026  0.865562

First, define the DoubleMLData object.

[3]:
data_dml_base = dml.DoubleMLData(
    data,
    y_col='y',
    d_cols='d'
)

Next, define the learners for the nuisance functions and fit the IRM Model. Remark that linear learners would usually be optimal due to the data generating process.

[4]:
# First stage estimation
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
randomForest_reg = RandomForestRegressor(n_estimators=500)
randomForest_class = RandomForestClassifier(n_estimators=500)

np.random.seed(42)

dml_irm = dml.DoubleMLIRM(data_dml_base,
                          ml_g=randomForest_reg,
                          ml_m=randomForest_class,
                          trimming_threshold=0.05,
                          n_folds=5)
print("Training IRM Model")
dml_irm.fit()

print(dml_irm.summary)
Training IRM Model
       coef  std err           t  P>|t|     2.5 %    97.5 %
d  4.475569   0.0408  109.695045    0.0  4.395603  4.555536

To estimate the CATE, we rely on the best-linear-predictor of the linear score as in Semenova et al. (2021) To approximate the target function \(\theta_0(x)\) with a linear form, we have to define a data frame of basis functions. Here, we rely on patsy to construct a suitable basis of B-splines.

[5]:
import patsy
design_matrix = patsy.dmatrix("bs(x, df=5, degree=2)", {"x": data["X_0"]})
spline_basis = pd.DataFrame(design_matrix)

To estimate the parameters to calculate the CATE estimate call the cate() method and supply the dataframe of basis elements.

[6]:
cate = dml_irm.cate(spline_basis)
print(cate)
================== DoubleMLBLP Object ==================

------------------ Fit summary ------------------
       coef   std err          t          P>|t|    [0.025    0.975]
0  0.691423  0.160438   4.309605   1.715075e-05  0.376780  1.006066
1  2.303007  0.267099   8.622301   1.314071e-17  1.779185  2.826829
2  4.904315  0.171709  28.561819  1.047375e-150  4.567568  5.241063
3  4.755688  0.205656  23.124465  5.235501e-105  4.352365  5.159011
4  3.745881  0.208922  17.929552   9.128273e-67  3.336153  4.155610
5  4.314341  0.224546  19.213635   1.278303e-75  3.873972  4.754710

To obtain the confidence intervals for the CATE, we have to call the confint() method and a supply a dataframe of basis elements. This could be the same basis as for fitting the CATE model or a new basis to e.g. evaluate the CATE model on a grid. Here, we will evaluate the CATE on a grid from 0.1 to 0.9 to plot the final results. Further, we construct uniform confidence intervals by setting the option joint and providing a number of bootstrap repetitions n_rep_boot.

[7]:
new_data = {"x": np.linspace(0.1, 0.9, 100)}
spline_grid = pd.DataFrame(patsy.build_design_matrices([design_matrix.design_info], new_data)[0])
df_cate = cate.confint(spline_grid, level=0.95, joint=True, n_rep_boot=2000)
print(df_cate)
       2.5 %    effect    97.5 %
0   2.070552  2.333655  2.596758
1   2.185984  2.453279  2.720573
2   2.298076  2.570936  2.843796
3   2.407558  2.686627  2.965696
4   2.515031  2.800351  3.085671
..       ...       ...       ...
95  4.417640  4.704814  4.991988
96  4.424292  4.705354  4.986417
97  4.433750  4.708235  4.982720
98  4.445476  4.713457  4.981438
99  4.458784  4.721018  4.983253

[100 rows x 3 columns]

Finally, we can plot our results and compare them with the true effect.

[8]:
from matplotlib import pyplot as plt
plt.rcParams['figure.figsize'] = 10., 7.5

df_cate['x'] = new_data['x']
df_cate['true_effect'] = treatment_effect(new_data["x"].reshape(-1, 1))
fig, ax = plt.subplots()
ax.plot(df_cate['x'],df_cate['effect'], label='Estimated Effect')
ax.plot(df_cate['x'],df_cate['true_effect'], color="green", label='True Effect')
ax.fill_between(df_cate['x'], df_cate['2.5 %'], df_cate['97.5 %'], color='b', alpha=.3, label='Confidence Interval')

plt.legend()
plt.title('CATE')
plt.xlabel('x')
_ =  plt.ylabel('Effect and 95%-CI')
../_images/examples_py_double_ml_cate_17_0.png

If the effect is not one-dimensional, the estimate still corresponds to the projection of the true effect on the basis functions.

Two-Dimensional Example#

It is also possible to estimate multi-dimensional conditional effects. We will use a similar data generating process but now the effect depends on the first two covariates \(X_0\) and \(X_1\) and takes the following form

\[\theta_0(X) = \exp(2X_0) + 3\sin(4X_1).\]

With the argument n_x=2 we can specify set the effect to be two-dimensional.

[9]:
np.random.seed(42)
data_dict = make_heterogeneous_data(
    n_obs=5000,
    p=10,
    support_size=5,
    n_x=2,
    binary_treatment=True,
)
treatment_effect = data_dict['treatment_effect']
data = data_dict['data']
print(data.head())
          y    d       X_0       X_1       X_2       X_3       X_4       X_5  \
0  1.286203  1.0  0.014080  0.006958  0.240127  0.100807  0.260211  0.177043
1  0.416899  1.0  0.152148  0.912230  0.892796  0.653901  0.672234  0.005339
2  2.087634  1.0  0.344787  0.893649  0.291517  0.562712  0.099731  0.921956
3  7.508433  1.0  0.619351  0.232134  0.000943  0.757151  0.985207  0.809913
4  0.567695  0.0  0.477130  0.447624  0.775191  0.526769  0.316717  0.258158

        X_6       X_7       X_8       X_9
0  0.028520  0.909304  0.008223  0.736082
1  0.984872  0.877833  0.895106  0.659245
2  0.140770  0.224897  0.558134  0.764093
3  0.460207  0.903767  0.409848  0.524934
4  0.037747  0.583195  0.229961  0.148134

As univariate example estimate the IRM Model.

[10]:
data_dml_base = dml.DoubleMLData(
    data,
    y_col='y',
    d_cols='d'
)
[11]:
# First stage estimation
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
randomForest_reg = RandomForestRegressor(n_estimators=500)
randomForest_class = RandomForestClassifier(n_estimators=500)

np.random.seed(42)

dml_irm = dml.DoubleMLIRM(data_dml_base,
                          ml_g=randomForest_reg,
                          ml_m=randomForest_class,
                          trimming_threshold=0.05,
                          n_folds=5)
print("Training IRM Model")
dml_irm.fit()

print(dml_irm.summary)
Training IRM Model
       coef   std err           t  P>|t|     2.5 %    97.5 %
d  4.547052  0.038845  117.057529    0.0  4.470918  4.623186

As above, we will rely on the patsy package to construct the basis elements. In the two-dimensional case, we will construct a tensor product of B-splines (for more information see here).

[12]:
design_matrix = patsy.dmatrix("te(bs(x_0, df=7, degree=3), bs(x_1, df=7, degree=3))", {"x_0": data["X_0"], "x_1": data["X_1"]})
spline_basis = pd.DataFrame(design_matrix)

cate = dml_irm.cate(spline_basis)
print(cate)
================== DoubleMLBLP Object ==================

------------------ Fit summary ------------------
         coef   std err          t          P>|t|    [0.025     0.975]
0    2.805766  0.105942  26.484022  1.169759e-144  2.598073   3.013459
1   -1.115489  0.865279  -1.289167   1.974002e-01 -2.811818   0.580841
2    0.462857  0.856112   0.540650   5.887734e-01 -1.215502   2.141215
3    2.309252  0.755712   3.055729   2.257106e-03  0.827721   3.790783
4    0.575588  0.763215   0.754162   4.507876e-01 -0.920652   2.071827
5   -3.338149  0.960068  -3.476991   5.114339e-04 -5.220309  -1.455990
6   -4.994680  1.033218  -4.834103   1.377947e-06 -7.020245  -2.969116
7   -6.404342  1.003959  -6.379085   1.943864e-10 -8.372547  -4.436136
8   -2.644977  0.908615  -2.911000   3.618824e-03 -4.426265  -0.863689
9    3.753215  0.930412   4.033929   5.567809e-05  1.929195   5.577234
10  -0.018440  0.787391  -0.023419   9.813172e-01 -1.562076   1.525197
11   1.577967  0.827434   1.907062   5.657037e-02 -0.044169   3.200104
12  -1.071240  1.027323  -1.042749   2.971154e-01 -3.085250   0.942769
13  -2.289879  1.135389  -2.016822   4.376796e-02 -4.515746  -0.064012
14  -2.607495  1.040439  -2.506150   1.223720e-02 -4.647216  -0.567774
15   0.241651  0.762233   0.317030   7.512342e-01 -1.252663   1.735965
16   1.085470  0.774248   1.401967   1.609878e-01 -0.432399   2.603340
17   3.742340  0.665550   5.622929   1.980483e-08  2.437567   5.047113
18   1.384833  0.673582   2.055922   3.984237e-02  0.064313   2.705353
19  -1.505902  0.842439  -1.787550   7.390974e-02 -3.157455   0.145652
20  -2.315927  0.896752  -2.582572   9.835149e-03 -4.073960  -0.557895
21  -2.806172  0.779345  -3.600679   3.204979e-04 -4.334034  -1.278310
22   1.945878  0.772286   2.519633   1.177894e-02  0.431855   3.459902
23   2.189109  0.805957   2.716160   6.627347e-03  0.609075   3.769142
24   4.375128  0.693628   6.307597   3.080755e-10  3.015309   5.734947
25   2.134494  0.666861   3.200809   1.379084e-03  0.827151   3.441837
26   1.709543  0.876075   1.951366   5.106987e-02 -0.007953   3.427038
27  -1.232706  0.945411  -1.303883   1.923340e-01 -3.086130   0.620719
28  -1.378910  0.907697  -1.519131   1.287934e-01 -3.158399   0.400578
29   4.060585  0.983891   4.127070   3.734540e-05  2.131723   5.989447
30   4.063691  0.973256   4.175357   3.026318e-05  2.155678   5.971704
31   6.076531  0.840625   7.228586   5.625914e-13  4.428533   7.724529
32   3.840627  0.818586   4.691784   2.781283e-06  2.235836   5.445417
33   2.585250  1.056947   2.445961   1.448125e-02  0.513166   4.657334
34  -1.848004  1.182386  -1.562945   1.181294e-01 -4.166004   0.469996
35  -0.429441  1.156561  -0.371309   7.104233e-01 -2.696813   1.837930
36   7.038915  0.978549   7.193219   7.274986e-13  5.120526   8.957305
37   5.187746  1.021007   5.081010   3.892380e-07  3.186120   7.189372
38   7.151029  0.833112   8.583517   1.216029e-17  5.517761   8.784298
39   6.733851  0.913280   7.373259   1.941512e-13  4.943417   8.524285
40   2.381220  1.148943   2.072531   3.826764e-02  0.128782   4.633658
41   4.441293  1.263665   3.514613   4.443368e-04  1.963949   6.918636
42   1.731904  1.166072   1.485246   1.375424e-01 -0.554115   4.017923
43  10.068584  0.937852  10.735793   1.353723e-26  8.229978  11.907189
44   3.734608  1.133198   3.295637   9.889032e-04  1.513038   5.956177
45   9.197947  0.814406  11.294060   3.189538e-29  7.601351  10.794544
46   5.367100  0.825135   6.504509   8.563902e-11  3.749469   6.984731
47   5.925833  1.023154   5.791732   7.396516e-09  3.919998   7.931669
48   1.301554  1.049953   1.239631   2.151705e-01 -0.756818   3.359927
49   1.237434  1.039654   1.190237   2.340104e-01 -0.800748   3.275616

Finally, we create a new grid to evaluate and plot the effects.

[13]:
grid_size = 100
x_0 = np.linspace(0.1, 0.9, grid_size)
x_1 = np.linspace(0.1, 0.9, grid_size)
x_0, x_1 = np.meshgrid(x_0, x_1)

new_data = {"x_0": x_0.ravel(), "x_1": x_1.ravel()}
[14]:
spline_grid = pd.DataFrame(patsy.build_design_matrices([design_matrix.design_info], new_data)[0])
df_cate = cate.confint(spline_grid, joint=True, n_rep_boot=2000)
print(df_cate)
         2.5 %    effect    97.5 %
0     1.671843  2.379763  3.087683
1     1.681670  2.367085  3.052500
2     1.698653  2.357300  3.015947
3     1.720696  2.350332  2.979968
4     1.745574  2.346106  2.946637
...        ...       ...       ...
9995  3.716422  4.506686  5.296950
9996  3.831783  4.661438  5.491093
9997  3.939298  4.811754  5.684210
9998  4.041979  4.955766  5.869554
9999  4.142442  5.091607  6.040771

[10000 rows x 3 columns]
[15]:
import plotly.graph_objects as go

grid_array = np.array(list(zip(x_0.ravel(), x_1.ravel())))
true_effect = treatment_effect(grid_array).reshape(x_0.shape)
effect = np.asarray(df_cate['effect']).reshape(x_0.shape)
lower_bound = np.asarray(df_cate['2.5 %']).reshape(x_0.shape)
upper_bound = np.asarray(df_cate['97.5 %']).reshape(x_0.shape)

fig = go.Figure(data=[
    go.Surface(x=x_0,
               y=x_1,
               z=true_effect),
    go.Surface(x=x_0,
               y=x_1,
               z=upper_bound, showscale=False, opacity=0.4,colorscale='purp'),
    go.Surface(x=x_0,
               y=x_1,
               z=lower_bound, showscale=False, opacity=0.4,colorscale='purp'),
])
fig.update_traces(contours_z=dict(show=True, usecolormap=True,
                                  highlightcolor="limegreen", project_z=True))

fig.update_layout(scene = dict(
                    xaxis_title='X_0',
                    yaxis_title='X_1',
                    zaxis_title='Effect'),
                    width=700,
                    margin=dict(r=20, b=10, l=10, t=10))

fig.show()